package com.laboratory.websocketlearning.handler;

import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

@Component
public class ChatHandler extends TextWebSocketHandler {

    /**
     * 所有session的映射
     */
    private final Map<String, WebSocketSession> sessions = new ConcurrentHashMap<>();

    /**
     * 用户与session的映射
     */
    private final Map<String, WebSocketSession> userSessions = new ConcurrentHashMap<>();

    /**
     * 房间与用户映射
     */
    private final Map<String, Set<String>> rooms = new ConcurrentHashMap<>();

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        sessions.put(session.getId(), session);
        String user = (String) session.getAttributes().getOrDefault("user", session.getId());
        userSessions.put(user, session);
        session.sendMessage(new TextMessage("WELCOME " + user));
        System.out.println("CONNECT " + user);
    }

    /**
     * 约定文本协议（简版）：
     * 1) 广播：        "ALL:消息内容"
     * 2) 加入房间：     "JOIN:room-1"
     * 3) 房间发消息：   "ROOM:room-1:消息内容"
     * 4) 点对点：      "TO:userB:消息内容"
     *
     * @param session
     * @param message
     * @throws Exception
     */
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        System.out.println(message.getPayload());
        String payload = message.getPayload();
        String user = (String) session.getAttributes().getOrDefault("user", session.getId());
        if (payload.startsWith("ALL:")) {
            String msg = "[" + user + "] " + payload.substring(4);
            broadcastAll(msg);
            return;
        }
        if (payload.startsWith("JOIN:")) {
            String roomId = payload.substring(5).trim();
            rooms.computeIfAbsent(roomId, k -> ConcurrentHashMap.newKeySet()).add(session.getId());
            session.sendMessage(new TextMessage("JOINED " + roomId));
            return;
        }
        if (payload.startsWith("ROOM:")) {
            String[] parts = payload.split(":", 3);
            if (parts.length == 3) {
                String roomId = parts[1];
                String content = parts[2];
                String msg = "[" + roomId + "] [" + user + "] " + content;
                broadcastRoom(roomId, msg);
            }
            return;
        }
        if (payload.startsWith("TO:")) {
            String[] parts = payload.split(":", 3);
            if (parts.length == 3) {
                String toUser = parts[1];
                String content = parts[2];
                sendToUser(toUser, "[" + user + " -> you]" + content);
            }
            return;
        }
        session.sendMessage(new TextMessage("ECHO: " + payload));
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        closeAndRemove(session);
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        closeAndRemove(session);
    }

    /**
     * 广播
     *
     * @param message
     */
    private void broadcastAll(String message) {
        for (WebSocketSession session : sessions.values()) {
            if (session.isOpen()) {
                safeSend(session, message);
            }
        }
    }

    /**
     * 广播房间
     *
     * @param roomId
     * @param message
     */
    private void broadcastRoom(String roomId, String message) {
        Set<String> members = rooms.get(roomId);
        if (members == null) {
            return;
        }
        for (String member : members) {
            WebSocketSession webSocketSession = sessions.get(member);
            if (webSocketSession != null && webSocketSession.isOpen()) {
                safeSend(webSocketSession, message);
            }
        }
    }

    /**
     * 点对点发送
     *
     * @param toUser
     * @param message
     */
    private void sendToUser(String toUser, String message) {
        WebSocketSession webSocketSession = userSessions.get(toUser);
        if (webSocketSession != null && webSocketSession.isOpen()) {
            safeSend(webSocketSession, message);
        }
    }

    /**
     * 关闭并移除session
     *
     * @param session
     */
    private void closeAndRemove(WebSocketSession session) {
        sessions.remove(session.getId());
        String user = (String) session.getAttributes().get("user");
        if (user != null) {
            userSessions.compute(user, (k, v) -> v == session ? null : v);
        } else {
            userSessions.values().removeIf(v -> v == session);
        }
        rooms.values().forEach(members -> members.remove(session.getId()));
        rooms.entrySet().removeIf(entry -> entry.getValue().isEmpty());
    }

    /**
     * 安全发送
     *
     * @param session
     * @param message
     */
    private void safeSend(WebSocketSession session, String message) {
        try {
            session.sendMessage(new TextMessage(message));
        } catch (Exception e) {
            System.out.println(e.getMessage());
        }
    }

}
